
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer

class DocumentLoader:
    """
    用于从PDF文件中提取文本的类。
    """

    def __init__(self, file_path, min_line_length=1):
        """
        初始化DocumentLoader。

        参数:
            min_line_length (int): 每行的最小长度，小于该长度的行会被忽略。默认为1。
        """
        self.min_line_length = min_line_length
        self.file_path = file_path

    def load_and_split(self, page_numbers=None):
        """
        从PDF文件中提取文本（按指定页码）。

        参数:
            page_numbers (list, optional): 需要提取的页码范围。默认为None，表示提取所有页。

        返回:
            list: 提取的段落列表，每个段落是一个字符串。
        """
        paragraphs = []  # 用于存储提取的段落
        buffer = ''  # 用于临时存储当前正在构建的段落
        full_text = ''  # 用于存储从PDF中提取的全部文本

        # 提取全部文本
        for i, page_layout in enumerate(extract_pages(self.file_path)):
            # 如果指定了页码范围，跳过范围外的页
            if page_numbers is not None and i not in page_numbers:
                continue

            # 遍历页面中的每个元素
            for element in page_layout:
                # 检查元素是否为文本容器（LTTextContainer）
                if isinstance(element, LTTextContainer):
                    # 将文本内容添加到full_text中，并添加换行符
                    full_text += element.get_text() + '\n'

        # 将提取的文本按行分割
        lines = full_text.split('\n')

        # 处理每一行文本
        for text in lines:
            # 如果当前行的长度大于等于最小行长度要求
            if len(text) >= self.min_line_length:
                # 如果行末没有连字符，直接添加到buffer中；否则去掉连字符并拼接
                buffer += (' ' + text) if not text.endswith('-') else text.strip('-')
            # 如果当前行长度不足，且buffer不为空，表示一个段落结束
            elif buffer:
                paragraphs.append(buffer)  # 将buffer中的内容作为一个段落添加到paragraphs中
                buffer = ''  # 清空buffer，准备构建下一个段落

        # 如果循环结束后buffer中仍有内容，将其作为最后一个段落添加到paragraphs中
        if buffer:
            paragraphs.append(buffer)

        # 返回提取的所有段落
        return paragraphs


#############################
# 创建DocumentLoader实例，设置最小行长度为5
# loader = DocumentLoader(min_line_length=5)
#
# # 提取PDF中第1页和第3页的文本
# paragraphs = loader.extract_text_from_pdf("example.pdf", page_numbers=[0, 2])
#
# # 打印提取的段落
# for para in paragraphs:
#     print(para)